#include "stdafx.h"

#include "RUdpSocket.h"
#include "ClientList.h"
#include "RUdpPacket.h"

RUdpSocket::RUdpSocket()
: mCount(1)
, mProcessing(0)
, mLastRcvdSeqId(0)
, mLastSentSeqId(0)
, mSocket(0)
, mCleanupTimer(0)
, mBacklogCount(0)
, mStatus(Disconnected)
{
  CreateSocket();
}

RUdpSocket::RUdpSocket(const RUdpSocket& s, const sockaddr_in& remoteEndPoint, const sockaddr_in& localEndPoint)
: mCount(1)
, mProcessing(0)
, mLastRcvdSeqId(0)
, mLastSentSeqId(0)
, mSocket(s.GetSocket())
, mRemoteEndpoint(remoteEndPoint)
, mLocalEndpoint(localEndPoint)
, mCleanupTimer(0)
, mBacklogCount(0)
, mStatus(Disconnected)
{
  mOrderedPackets.resize(16);

  SetTimeStamp(timeGetTime());

  mHash = HashEndPoint();
}

RUdpSocket::~RUdpSocket()
{
  mCleanupTimer = 0;
}

int RUdpSocket::Connect(char* address, short port)
{
  memset((void *)&mRemoteEndpoint, 0, sizeof(struct sockaddr_in));

  mRemoteEndpoint.sin_family = AF_INET;
  mRemoteEndpoint.sin_addr.s_addr = inet_addr(address);
  mRemoteEndpoint.sin_port = htons(port);

  int status = connect(mSocket, (SOCKADDR*) &mRemoteEndpoint, sizeof(mRemoteEndpoint));

  if (status == 0)
  {
    RUdpPacket packet;
    packet.mPacketType = PT_SYN;
    packet.mSeqId = rand();
    SendTo(packet);

    std::vector<char> buffer(64);
    size_t bytes = Receive(buffer);

    if (bytes > 0)
    {
      RUdpPacket* packet = new RUdpPacket;
      packet->FromBuffer(&buffer[0], bytes);

      ProcessReceivedPacket(packet);
    }
    else
      status = -1;
  }

  return status;
}

UINT64 RUdpSocket::HashEndPoint()
{
  UINT64 hash = mRemoteEndpoint.sin_addr.s_addr;

  hash <<= 32;
  hash += (mRemoteEndpoint.sin_port << 16);
  hash += mLocalEndpoint.sin_port;

  return hash;
}

SOCKET RUdpSocket::CreateSocket()
{
  mSocket = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

  return mSocket;
}

size_t RUdpSocket::Receive(std::vector<char>& buffer)
{
  int sockAddrLen = sizeof(sockaddr_in);

  return recvfrom(mSocket, &buffer[0], buffer.size(), 0, (SOCKADDR*) &mRemoteEndpoint, &sockAddrLen);
}

size_t RUdpSocket::Receive(char* buffer, int len)
{
  int sockAddrLen = sizeof(sockaddr_in);

  return recvfrom(mSocket, buffer, len, 0, (SOCKADDR*) &mRemoteEndpoint, &sockAddrLen);
}

bool RUdpSocket::Bind(char* address, short port)
{
  memset((void *)&mLocalEndpoint, 0, sizeof(struct sockaddr_in));

  mLocalEndpoint.sin_family = AF_INET;
  mLocalEndpoint.sin_addr.s_addr = inet_addr(address);
  mLocalEndpoint.sin_port = htons(port);

  if (bind( mSocket, (SOCKADDR*) &mLocalEndpoint, sizeof(sockaddr_in) ) == SOCKET_ERROR)
    return false;

  return true;
}

void RUdpSocket::KillCleanupTimer()
{
  if ( mCleanupTimer )
  {
    SetThreadpoolTimer(mCleanupTimer, NULL, 0, 0);
    CloseThreadpoolTimer(mCleanupTimer);
  }
}

int RUdpSocket::SendTo(const RUdpPacket& packet)
{
  static char buffer[65 * 1024];

  int packetLen = packet.ToBuffer(buffer, sizeof(buffer));

  int bytes = sendto(mSocket, &buffer[0], packetLen, 0, (sockaddr*) &mRemoteEndpoint, sizeof(sockaddr_in));

  if (bytes == packetLen)
    return packet.mPayloadLen;

  return bytes;
}

int RUdpSocket::SendTo(const char* data, size_t len)
{
  RUdpPacket packet;
  packet.mPacketType = PT_MSG;

  return SendTo(packet);
}

void RUdpSocket::ProcessOutOfOrderPackets()
{
  int index = 0;
  int backlogCount = mBacklogCount;

  while ( index < backlogCount )
  {
    RUdpPacket* packet = mOrderedPackets[index];

    if (packet != 0)
    {
      ProcessReceivedPacket(packet, true);
      --mBacklogCount;
      ++index;
    }
    else
    {
      break;
    }
  }

  // Shift vector (ringbuffer would be better)
  mOrderedPackets.erase(mOrderedPackets.begin(), mOrderedPackets.begin() + index);
  mOrderedPackets.resize(16);
}

void RUdpSocket::ProcessMessagePacket(RUdpPacket* packet, bool outOfOrder)
{
  // In-order packet received
  if (packet->mSeqId == mLastRcvdSeqId + 1)
  {
    std::cout << "Processing packet... Status : " << packet->mPacketType << " : Sequence (" << packet->mSeqId << ")" << std::endl;

    mLastRcvdSeqId = packet->mSeqId;

    // Send ACK
    RUdpPacket ack;
    ack.mPacketType = PT_ACK;
    ack.mAckId = packet->mSeqId + 1;

    SendTo(ack);

    // TODO: Process packet

    // Don't need the packet anymore
    delete packet;

    // Prevent calling this again while processing out-of-order packets
    if (!outOfOrder)
      ProcessOutOfOrderPackets();
  }
  // Processing an out-of-order packet.
  else if (packet->mSeqId > mLastRcvdSeqId + 1)
  {
    size_t index = packet->mSeqId - mLastRcvdSeqId - 2;

    if (index < 0 || index >= mOrderedPackets.size())
    {
      //data = ((__int64) 0 << 32) + mLastRcvdSeqId + 1;
      //SendTo((char*)&data, sizeof(data));
    }
    else if (mOrderedPackets[index] == 0)
    {
      std::cout << "Out of order packet : Sequence (" << packet->mSeqId << ")" << std::endl;

      mOrderedPackets[index] = packet;
      ++mBacklogCount;
    }
  }
  // Packet that is below our sequence number. Will just ignore
  // in the future, but client needs a response right now so just send one.
  else
  {
    delete packet;
  }
}

void RUdpSocket::ProcessReceivedPacket(RUdpPacket* packet, bool outOfOrder)
{
  // Mark the Socket timestamp with the latest segment
  SetTimeStamp( packet->mTimeStamp );

  switch (packet->mPacketType)
  {
  case PT_SYN:
    {
      mLastRcvdSeqId = packet->mSeqId;
      mLastSentSeqId = 1;

      RUdpPacket ack;
      ack.mPacketType = PT_SYNACK;
      ack.mAckId = packet->mSeqId + 1;
      ack.mSeqId = rand();

      SendTo(ack);

      mStatus = Connecting;

      delete packet;
    }
    break;

  case PT_ACK:
    {
      mLastRcvdSeqId = packet->mSeqId;

      if (mStatus == Connecting)
        mStatus = Connected;
      
      delete packet;
    }
    break;

  case PT_SYNACK:
    {
      mLastRcvdSeqId = packet->mSeqId;
      mLastSentSeqId++;

      RUdpPacket ack;
      ack.mPacketType = PT_ACK;
      ack.mAckId = packet->mSeqId + 1;
      ack.mSeqId = packet->mAckId;

      SendTo(ack);

      mStatus = Connected;

      delete packet;
    }
    break;

  case PT_MSG:
    {
      ProcessMessagePacket(packet, outOfOrder);
    }
    break;
  }
}

// Should handle ordering and reliability here
void RUdpSocket::OnPacketReceive()
{
  long count = InterlockedIncrement(&mProcessing);

  if ( count == 1 )
  {
    while (RUdpPacket* packet = mPackets.Pop())
    {
      ProcessReceivedPacket( packet );
    }
  }

  InterlockedDecrement(&mProcessing);
}

void RUdpSocket::Push(RUdpPacket* packet)
{
  mPackets.Push(packet);
}

// Crude version of a smart pointer. It's all we need for this example.
long RUdpSocket::AddRef()
{
  return InterlockedIncrement(&mCount);
}

void RUdpSocket::Release()
{
  if (InterlockedDecrement(&mCount) == 0)
    delete this;
}

